import json
import os
import sqlite3
from typing import TypedDict, Annotated, List

import dotenv
from langchain_core.messages import BaseMessage, AIMessage, ToolMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langchain_core.runnables import Runnable
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver

# 导入所有可用的工具和工具检索器
from tools import all_available_tools
from tool_retriever import get_tool_retriever

dotenv.load_dotenv()


# 1. 定义Graph的状态
class AgentState(TypedDict):
    messages: Annotated[List[BaseMessage], lambda x, y: x + y]
    # 【新增】一个字段来存储当前 Agent 可用的工具
    available_tools: List[Runnable]


# 2. 定义节点和逻辑
# 不再绑定所有工具
model = ChatOpenAI(
    model="deepseek-chat",
    base_url=os.getenv("DS_BASE"),
    api_key=os.getenv("DS_API_KEY"),
    temperature=0
)
# 初始化工具检索器
tool_retriever = get_tool_retriever(all_available_tools)


def agent_node(state: AgentState):
    """
    Agent决策节点。根据当前状态中的 available_tools 来决定下一步行动。
    """
    print("---NODE: Agent is thinking...---")
    # 从状态中获取当前可用的工具
    tools_for_this_turn = state.get("available_tools", [])
    # 动态绑定工具，而不是一次性绑定所有工具
    model_with_tools = model.bind_tools(tools_for_this_turn)

    response = model_with_tools.invoke(state["messages"])
    print(f"AI thinking ret: {response}")
    return {"messages": [response]}


def tool_retriever_node(state: AgentState):
    """
    工具检索节点。根据用户的最新输入，检索相关工具。
    """
    print("---NODE: Retrieving tools...---")
    # 获取最新的用户输入
    last_message = state["messages"][-1]
    if not isinstance(last_message, HumanMessage):
        # 如果不是用户消息，则使用之前 AI 的消息内容进行检索
        last_message = next((msg for msg in reversed(state["messages"]) if isinstance(msg, HumanMessage)), None)
        if not last_message:
            # 找不到用户消息，无法检索，返回所有工具作为备选
            return {"available_tools": all_available_tools}

    # 使用 RAG 检索工具
    retrieved_tools = tool_retriever.invoke(last_message.content)
    print(f"---Retrieved {len(retrieved_tools)} tools.---")

    return {"available_tools": retrieved_tools}


def tool_executor_node(state: AgentState):
    """
    工具执行节点。手动执行 Agent 提出的工具调用。
    """
    print(f"---NODE: Executing tools...---")
    last_message = state["messages"][-1]
    tool_calls = last_message.tool_calls

    responses = []
    for tool_call in tool_calls:
        tool_name = tool_call["name"]
        tool_args = tool_call["args"]

        # 从状态中找到对应的工具并执行
        # 这要求 tool_retriever_node 已经将工具放在了 state['available_tools'] 中
        tool_func = next((t for t in state["available_tools"] if t.name == tool_name), None)

        if not tool_func:
            responses.append(f"工具 '{tool_name}' 未找到或未被检索到。")
            continue

        try:
            result = tool_func.invoke(tool_args)
            responses.append(str(result))
        except Exception as e:
            responses.append(f"执行工具 '{tool_name}' 失败: {str(e)}")

    tool_messages = [
        ToolMessage(content=res, tool_call_id=call.id)
        for res, call in zip(responses, last_message.tool_calls)
    ]
    return {"messages": tool_messages}


def should_continue(state: AgentState) -> str:
    """条件路由：判断是否需要调用工具。"""
    print("---ROUTER: Deciding next step---")
    # 如果 Agent 提出工具调用，则进入工具执行流程
    if state["messages"][-1].tool_calls:
        print("Go execute tools.")
        return "continue"
    print("Go end.")
    return "end"


def should_retrieve_tools(state: AgentState) -> str:
    """条件路由：判断是否需要检索工具。"""
    print("---ROUTER: Deciding whether to retrieve tools---")
    # 如果状态中还没有 available_tools，说明是新对话，需要先检索工具
    if not state.get("available_tools"):
        print("Initial input, need to retrieve tools.")
        return "retrieve"
    print("Tools already available, go directly to agent.")
    return "skip"


# 3. 构建并编译Graph
workflow = StateGraph(AgentState)

# 【新增】工具检索节点
workflow.add_node("tool_retriever", tool_retriever_node)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_executor_node)

# 新的入口点和路由
workflow.set_entry_point("tool_retriever")
workflow.add_edge("tool_retriever", "agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    {"continue": "tools", "end": END},
)
workflow.add_edge("tools", "agent")

# 设置检查点
conn = sqlite3.connect("checkpoints.sqlite",
                       check_same_thread=False,
                       isolation_level=None,
                       timeout=30,
                       detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES)
memory = SqliteSaver(conn)

# 编译Graph，并设置中断条件
app_graph = workflow.compile(
    checkpointer=memory,
    interrupt_before=["tools"]  # 依然在执行工具前中断
)
